import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from sklearn import tree

df = pd.read_csv('合并RID(填充缺失值+去typeNaN).csv', encoding='gbk')
Xcols = df.columns.tolist()[4:]
morethan03 = ['PIB', 'AV45', 'FBB', 'DIGITSCOR', 'DIGITSCOR_bl', 'MriType2', 'MOCA_bl', 'EcogPtMem_bl', 'EcogPtLang_bl', 'EcogPtVisspat_bl', 'EcogPtPlan_bl', 'EcogPtOrgan_bl', 'EcogPtDivatt_bl', 'EcogPtTotal_bl', 'EcogSPMem_bl', 'EcogSPLang_bl', 'EcogSPVisspat_bl', 'EcogSPPlan_bl', 'EcogSPOrgan_bl', 'EcogSPDivatt_bl', 'EcogSPTotal_bl', 'FDG_bl', 'PIB_bl', 'AV45_bl', 'FBB_bl']
nowNanCols = ['MOCA','EcogPtMem','EcogPtLang','EcogPtVisspat','EcogPtPlan','EcogPtOrgan','EcogPtDivatt','EcogPtTotal', 'update_stamp', 'TYPE2', 'EXAMDATE_bl', 'Years_bl', 'Month_bl', 'Month', 'M']
Xcols = list(set(Xcols) - set(morethan03) - set(nowNanCols))
# 去掉有nan的列
realXcols = []
for i in Xcols:
    if df[i].isnull().any():
        print(i, 'has NaN')
    else:
        realXcols.append(i)
print(realXcols)

clf = DecisionTreeClassifier(max_depth = 5,
                             random_state = 0)
clf.fit(df[realXcols], df['TYPE'])
plt.figure(figsize=(30,5))
tree.plot_tree(clf, fontsize=7, feature_names=realXcols, class_names=['AD', 'CN', 'EMCI', 'LMCI', 'SMC'])
plt.savefig('tree_high_dpi', dpi=100)

# 验证结果
'''
rightNum = 0
rightData = df['TYPE'].tolist()
for i in range(len(ret)):
    if ret[i] == rightData[i]:
        rightNum += 1
print(rightNum, '/', len(ret), rightNum/len(ret))
'''